# coding:utf8

import os

from utils.content.file_reader import ReadFiles
from utils.storage.vector_store import VectorStore
from utils.embedding.dashscope_embedding import DashscopeEmbedding
from utils.llm.deepseek_model import DeepSeekModel

def run_mini_rag(question: str, knowledge_base_path: str, k: int = 1) -> str:
    """
    运行一个简化版的RAG项目。
    
    :param question: 用户提出的问题
    :param knowledge_base_path: 知识库的路径，包含文档的文件夹路径
    :param k: 返回与问题最相关的k个文档片段，默认为1
    :return: 返回模型生成的回答
    """
    # 1. 加载并切分文档
    docs = ReadFiles(knowledge_base_path).get_content(max_token_len=600, cover_content=150)
    vector = VectorStore(docs)
    print(f"以将文档切分为 {len(docs)} 份。")

    # 2. 使用 Dashscope Embedding 模型进行向量化
    embedding = DashscopeEmbedding()
    vector.get_vector(EmbeddingModel=embedding)
    print("向量化完成。")

    # 3. 将向量和文档保存到本地（可选）
    vector.persist(path='storage')
    print("向量存储完成。")

    # 4. 在数据库中检索最相关的文档片段
    content = vector.query(question, EmbeddingModel=embedding, k=k)[0]
    print(f"匹配问题上下文: {content}")

    # 5. 使用 DeepSeek 生成答案
    chat =  DeepSeekModel()
    answer = chat.chat(question, [], '')
    
    return answer

# test
question = '如何在Gerrit项目中导入LDAP证书？'
knowledge_base_path = './data'
answer = run_mini_rag(question, knowledge_base_path)
print(f"模型作答: {answer}")